/*
ionet
*/
package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/songgao/water"
	"github.com/songgao/water/waterutil"
)

type config struct {
	Devtype    string
	Devnetwork string
	Localaddr  string
	Peeraddrs  string
	timestamp  bool
	queue      int
}

type Table struct {
	Ipmap   *sync.Map
	Peers   []string
	tuntap  string
	localIP net.IP
	log     func(a ...any)
}

func (t *Table) AddElms(peer ...string) {
	ps := make(map[string]bool, 0)
	tmp := make([]string, 0)
	for _, v := range t.Peers {
		ps[v] = true
	}
	for _, v := range peer {
		ps[v] = true
	}
	for k := range ps {
		tmp = append(tmp, k)
	}
	t.Peers = tmp
}

func (t *Table) DeleteElms(peer ...string) {
	ps := make(map[string]bool, 0)
	tmp := make([]string, 0)
	for _, v := range t.Peers {
		ps[v] = true
	}
	for _, v := range peer {
		delete(ps, v)
	}
	t.Peers = tmp
}

func main() {
	var Conf = &config{}
	flag.IntVar(&Conf.queue, "q", 20, "queue of packet or frame")
	flag.BoolVar(&Conf.timestamp, "t", false, "log timestamp")
	flag.StringVar(&Conf.Devtype, "i", "tun", "type of tuntap device[tap|tun] ")
	flag.StringVar(&Conf.Devnetwork, "net", "10.0.0.1/32", "ip address of tun/tap")
	flag.StringVar(&Conf.Peeraddrs, "p", "", "peer addresses used [,] split")
	flag.StringVar(&Conf.Localaddr, "l", "0.0.0.0:8001", "local  address")
	flag.Parse()

	TransUdp(Conf, createtuntap(Conf), nil)
}

func createtuntap(conf *config) *water.Interface {
	fmt.Println(SetColor("1", "32", "["+os.Args[0]+" init]"))
	defer fmt.Println(SetColor("1", "32", "["+os.Args[0]+" start]"))
	// 选择tuntap类型，及设置配置
	config := water.Config{}
	switch conf.Devtype {
	case "tap":
		config.DeviceType = water.TAP
	default:
		config.DeviceType = water.TUN
	}
	// 根据配置生成接口
	ifce, err := water.New(config)
	if err != nil {
		log.Fatal(err)
	}
	if runtime.GOOS == "linux" {
		err = setupNetDevice("ip", "addr", "add", conf.Devnetwork, "dev", ifce.Name())
		if err != nil {
			fmt.Println("[", ifce.Name(), "]:设置网络地址失败：", err)
		} else {
			fmt.Println("[", ifce.Name(), "]:设置地址：", conf.Devnetwork)
		}
		err = setupNetDevice("ip", "link", "set", ifce.Name(), "mtu", "1400", "up")
		if err != nil {
			fmt.Println("[", ifce.Name(), "]:启动失败：", err)
		} else {
			fmt.Println("[", ifce.Name(), "]:启动成功")
		}
	}
	return ifce
}

func setupNetDevice(cmd ...string) error {
	cmdexec := exec.Command(cmd[0], cmd[1:]...)
	cmdexec.Stderr = os.Stderr
	cmdexec.Stdout = os.Stdout
	err := cmdexec.Run()
	if err != nil {
		return err
	}
	return nil
}

type Data struct {
	buf  []byte
	addr string
	n    int
}

func NewPacket() *Data {
	return &Data{buf: make([]byte, 1500)}
}

/*
udp
*/
//需要自定义proto function返回字节切片和数据长度
func TransUdp(conf *config, ifio *water.Interface, proto func(encode bool, b []byte) (int, []byte)) {
	defer ifio.Close()

	//1.解析公网节点地址
	var tab = &Table{Peers: make([]string, 0), Ipmap: &sync.Map{}, log: func() func(a ...any) {
		if conf.timestamp {
			return func(a ...any) { fmt.Println(append([]any{time.Now().Format("2006-01-02 15:04:05.000")}, a...)...) }
		} else {
			return func(a ...any) { fmt.Println(a...) }
		}
	}()}
	if ifio.IsTAP() {
		tab.tuntap = "TAP"
	} else {
		tab.tuntap = "TUN"
	}
	tab.Peers = strings.Split(conf.Peeraddrs, ",")
	for _, peer := range tab.Peers {
		paddr, err := net.ResolveUDPAddr("udp", peer)
		if err != nil {
			tab.log(SetColor("3", "31", "parse udp address error:"), err)
			tab.DeleteElms(peer)
			continue
		} else if peer == "" {
			tab.DeleteElms(peer)
			continue
		}
		tab.log(SetColor("2", "32", "[Add peer]:"), "<==", paddr.String())
	}
	//2.监听udp
	locaddr, err := net.ResolveUDPAddr("udp", conf.Localaddr)
	if err != nil {
		tab.log(SetColor("3", "31", "parse udp address error:"), err)
		return
	}
	listener, err := net.ListenUDP("udp", locaddr)
	if err != nil {
		tab.log(SetColor("3", "31", "listen error:"), err)
		return
	}
	//3.解析本地址
	tab.localIP, _, err = net.ParseCIDR(conf.Devnetwork)
	if err != nil {
		tab.log(SetColor("3", "31", "parse local address error:"), err)
		return
	}
	defer listener.Close()
	fromdev := make(chan *Data, conf.queue)
	fromnet := make(chan *Data, conf.queue)
	todev := make(chan *Data, conf.queue)
	tonet := make(chan *Data, conf.queue)
	defer close(fromdev)
	defer close(fromnet)
	defer close(todev)
	defer close(tonet)
	// dev -> chan
	go func() {
		for {
			packet := NewPacket()
			packet.n, err = ifio.Read(packet.buf)
			if err != nil {
				tab.log(SetColor("3", "31", "resv device error:"), err)
				continue
			}
			fromdev <- packet
		}
	}()
	// net -> chan
	go func() {
		for {
			packet := NewPacket()
			n, remotAddr, err := listener.ReadFromUDP(packet.buf)
			if err != nil {
				tab.log(SetColor("3", "31", "resv udp error:"), err)
				continue
			}
			packet.addr = remotAddr.String()
			packet.n = n
			//自定义协议
			if proto != nil {
				packet.n, packet.buf = proto(false, packet.buf)
			}
			fromnet <- packet
		}
	}()
	// chan -> net or device
	go func() {
		for {
			select {
			case packet := <-todev:
				ifio.Write(packet.buf[:packet.n])
			case packet := <-tonet:
				remoteaddr, err := net.ResolveUDPAddr("udp", packet.addr)
				if err != nil {
					tab.log(SetColor("3", "31", "parse udp address error:"), err)
					continue
				}
				if proto != nil {
					packet.n, packet.buf = proto(false, packet.buf)
				}
				listener.WriteToUDP(packet.buf[:packet.n], remoteaddr)
			}
		}
	}()
	// parse network packet
	ParseData(tab, fromdev, fromnet, todev, tonet)
}

// parse network packet
func ParseData(tab *Table, fromdev, fromnet <-chan *Data, todev, tonet chan<- *Data) {
	// keepalive udp conn
	go func() {
		if len(tab.Peers) > 0 {
			for {
				for _, peer := range tab.Peers {
					packet := NewPacket()
					packet.buf = tab.localIP.To4()
					packet.addr = peer
					packet.n = len(packet.buf)
					tonet <- packet
				}
				time.Sleep(7 * time.Second)
			}
		} else {
			return
		}
	}()
	switch tab.tuntap {
	case "TAP":
		for {
			select {
			case frame := <-fromnet:
				// parse frame from net
				/* keepalive*/
				if ip, ok := keepalive(frame.buf[:frame.n]); ok {
					_, ok := tab.Ipmap.Load(ip.String())
					if !ok {
						tab.log(SetColor("2", "32", "[Ipmap]add:"), ip.String(), "==>", frame.addr)
					}
					tab.Ipmap.Store(ip.String(), frame.addr)
					continue
				}
				dstMAC := waterutil.MACDestination(frame.buf[:frame.n])
				switch waterutil.MACEthertype(frame.buf[:frame.n]) {
				case waterutil.ARP:
					ARPsIP := ARPsrcIP(frame.buf[:frame.n])
					ARPdIP := ARPdstIP(frame.buf[:frame.n])
					if waterutil.IsBroadcast(dstMAC) {
						v, ok := tab.Ipmap.LoadOrStore(ARPsIP.String(), frame.addr)
						if !ok {
							tab.log(SetColor("2", "32", "[Ipmap]:add"), ARPsIP.String(), "==>", frame.addr)
						} else {
							if v.(string) != frame.addr {
								tab.Ipmap.Store(ARPsIP.String(), frame.addr)
							}
						}
						// 判断是否转发
						if ARPdIP.String() != tab.localIP.String() {
							v, ok := tab.Ipmap.Load(ARPdIP.String())
							if !ok {
								tab.log(SetColor("3", "31", "not found ip in map:"), ARPdIP.String())
								continue
							} else {
								frame.addr = v.(string)
								tonet <- frame
							}
						} else {
							todev <- frame
						}
					} else {
						v, ok := tab.Ipmap.LoadOrStore(ARPsIP.String(), frame.addr)
						if !ok {
							tab.log(SetColor("2", "32", "[Ipmap]:add"), ARPsIP.String(), "==>", frame.addr)
						} else {
							if v.(string) != frame.addr {
								tab.log(ARPsIP.String(), frame.addr)
							}
						}
						// 判断是否转发
						if ARPdIP.String() != tab.localIP.String() {
							v, ok := tab.Ipmap.Load(ARPdIP.String())
							if !ok {
								tab.log(SetColor("3", "31", "not found ip in map:"), ARPdIP.String())
								continue
							} else {
								frame.addr = v.(string)
								tonet <- frame
							}
						} else {
							todev <- frame
						}
					}
				case waterutil.IPv4:
					dstIP := waterutil.IPv4Destination(waterutil.MACPayload(frame.buf[:frame.n]))
					if dstIP.String() != tab.localIP.String() {
						v, ok := tab.Ipmap.Load(dstIP.String())
						if !ok {
							tab.log(SetColor("3", "31", "not found ip in map:"), dstIP.String())
							continue
						} else {
							frame.addr = v.(string)
							tonet <- frame
						}
					} else {
						todev <- frame
					}
				}
			case frame := <-fromdev:
				dstMAC := waterutil.MACDestination(frame.buf[:frame.n])
				switch waterutil.MACEthertype(frame.buf[:frame.n]) {
				case waterutil.ARP:
					ARPdIP := ARPdstIP(frame.buf[:frame.n])
					if waterutil.IsBroadcast(dstMAC) {
						v, ok := tab.Ipmap.Load(ARPdIP.String())
						if !ok {
							if len(tab.Peers) > 0 {
								for _, peer := range tab.Peers {
									tonet <- &Data{buf: frame.buf, addr: peer, n: frame.n}
								}
							}
						} else {
							frame.addr = v.(string)
							tonet <- frame
						}
					} else {
						v, ok := tab.Ipmap.Load(ARPdIP.String())
						if ok {

							frame.addr = v.(string)
							tonet <- frame
						} else {
							tab.log(SetColor("3", "31", "found Ipmap error"))
							continue
						}
					}
				case waterutil.IPv4:
					v, ok := tab.Ipmap.Load(waterutil.IPv4Destination(waterutil.MACPayload(frame.buf[:frame.n])).String())
					if !ok {
						tab.log(SetColor("3", "31", "not found ip in map:"), waterutil.IPv4Destination(waterutil.MACPayload(frame.buf[:frame.n])).String())
						continue
					} else {
						frame.addr = v.(string)
						tonet <- frame
					}
				}
			}
		}
	case "TUN":
		for {
			select {
			case packet := <-fromnet:
				// parse packet from net
				/* keepalive*/
				if ip, ok := keepalive(packet.buf[:packet.n]); ok {
					_, ok := tab.Ipmap.Load(ip.String())
					if !ok {
						tab.log(SetColor("2", "32", "[Ipmap]add:"), ip.String(), "==>", packet.addr)
					}
					tab.Ipmap.Store(ip.String(), packet.addr)
					continue
				}
				if waterutil.IsIPv4(packet.buf[:packet.n]) {
					srcIP := waterutil.IPv4Source(packet.buf[:packet.n])
					dstIP := waterutil.IPv4Destination(packet.buf[:packet.n])
					// 查ipmap是否有src ipmap
					v, ok := tab.Ipmap.LoadOrStore(srcIP.String(), packet.addr)
					if !ok {
						tab.log(SetColor("2", "32", "[Ipmap]:add"), srcIP.String(), "==>", packet.addr)
					} else {
						if v.(string) != packet.addr {
							tab.Ipmap.Store(srcIP.String(), packet.addr)
						}
					}
					// 判断是否转发
					if dstIP.String() != tab.localIP.String() {
						v, ok := tab.Ipmap.Load(dstIP.String())
						if !ok {
							tab.log(SetColor("3", "31", "not found ip in map:"), dstIP.String())
							continue
						} else {
							packet.addr = v.(string)
							tonet <- packet
						}
					} else {
						todev <- packet
					}
				}
			case packet := <-fromdev:
				// parse data from device
				dstIP := waterutil.IPv4Destination(packet.buf[:packet.n])
				v, ok := tab.Ipmap.Load(dstIP.String())
				if ok {
					packet.addr = v.(string)
					tonet <- packet
				} else {
					if len(tab.Peers) < 1 {
						tab.Ipmap.Range(func(key, value any) bool {
							tonet <- &Data{buf: packet.buf, addr: value.(string), n: packet.n}
							return true
						})
					} else {
						for _, peer := range tab.Peers {
							tonet <- &Data{buf: packet.buf, addr: peer, n: packet.n}
						}
					}
				}
			}
		}
	}
}

func keepalive(buf []byte) (net.IP, bool) {
	if len(buf) != 4 {
		return nil, false
	} else {
		return net.IPv4(buf[0], buf[1], buf[2], buf[3]), true
	}
}

// color tools
func SetColor(i, c, in string) string {
	return fmt.Sprint("\033[", i, ";", c, "m", in, "\033[0m")
}

// ARP tools
func ARPsrcIP(frame []byte) net.IP {
	return net.IPv4(frame[28], frame[29], frame[30], frame[31])
}
func ARPdstIP(frame []byte) net.IP {
	return net.IPv4(frame[38], frame[39], frame[40], frame[41])
}

//[255 255 255 255 255 255 34 210 66 114 118 138 8 6 0 1 8 0 6 4 0 1 34 210 66 114 118 138 10 0 0 4 0 0 0 0 0 0 10 0 0 9]
//                  dstmac                srcmac12131415161718192021                srcmac 28                   38
